from tensorflow import keras
from .config import *

def build_model():
    input_image = keras.Input(shape=(config.IMG_SIDE_LEN, config.IMG_SIDE_LEN, config.NUM_LOG_TERM*config.NUM_COLOR_CHANNEL))
    input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
    input_pv_mlp =  keras.Input(shape=(config.NUM_LOG_TERM,))
    if config.MODEL_SELECT == "CNN_LSTM":
        # CNN分支
        x = keras.layers.Conv2D(32, (3,3), activation='relu', padding='same')(input_image)
        x = keras.layers.MaxPooling2D(2)(x)
        x = keras.layers.Conv2D(64, (3,3), activation='relu', padding='same')(x)
        x = keras.layers.MaxPooling2D(2)(x)
        x = keras.layers.Flatten()(x)
        x = keras.layers.Dense(128, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        cnn_output = keras.layers.Dense(64, activation='relu')(x)
        # LSTM分支
        y = keras.layers.LSTM(64, return_sequences=True)(input_pv_lstm)
        y = keras.layers.Dropout(0.2)(y)
        y = keras.layers.LSTM(32, return_sequences=False)(y)
        y = keras.layers.Dropout(0.2)(y)
        lstm_output = keras.layers.Dense(32, activation='relu')(y)
        # 合并分支
        combined = keras.layers.concatenate([cnn_output, lstm_output])
        z = keras.layers.Dense(64, activation='relu')(combined)
        z = keras.layers.Dropout(0.2)(z)
        output = keras.layers.Dense(units=1)(z)
        model = keras.Model(inputs=[input_image, input_pv_lstm], outputs=output)
    elif config.MODEL_SELECT == "LSTM":
        x = keras.layers.LSTM(64, return_sequences=True)(input_pv_lstm)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.LSTM(64, return_sequences=True)(x)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.LSTM(32, return_sequences=False)(x)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.Dense(16, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        output = keras.layers.Dense(units=1)(x)
        model = keras.Model(inputs=input_pv_lstm, outputs=output)
    elif config.MODEL_SELECT == "MLP":
        # Fully connected layers
        x = keras.layers.Dense(64, activation='relu')(input_pv_mlp)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.Dense(32, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.Dense(16, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        # Regression to prediction target
        output = keras.layers.Dense(units=1)(x)
        # Construct the model
        model = keras.Model(inputs=input_pv_mlp, outputs=output)
    else:
        print("Selectd model is error.")
    return model 